use serde::{Deserialize, Serialize};
use std::fs;
use std::path::Path;
use std::sync::Arc;
use tokio::sync::Mutex;
use uuid::Uuid;

// Import the new agent types
use crate::agent_types::{AgentCategory, EnhancedAgentConfig, KnowledgeBaseConfig, VectorDbConfig};

// Define our own Message struct that's compatible with the API
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
    pub role: String,
    pub content: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentConfig {
    pub id: i64,
    pub uuid: String,
    pub name: String,
    pub provider: String,
    pub model: String,
    pub preamble: String,
    pub created_at: String,
    pub base_url: Option<String>,
    pub tools: Option<Vec<String>>,
    pub api_key: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct AgentsData {
    agents: Vec<EnhancedAgentConfig>,
}

#[derive(Debug, Clone)]
pub struct AgentManager {
    data_path: String,
    data: Arc<Mutex<AgentsData>>,
}

impl AgentManager {
    pub fn new<P: AsRef<Path>>(data_path: P) -> anyhow::Result<Self> {
        let data_path = data_path.as_ref().to_string_lossy().to_string();
        let manager = Self {
            data_path: data_path.clone(),
            data: Arc::new(Mutex::new(AgentsData { agents: vec![] })),
        };
        // Load data synchronously
        manager.load_data_sync()?;
        Ok(manager)
    }

    fn load_data_sync(&self) -> anyhow::Result<()> {
        if Path::new(&self.data_path).exists() {
            let content = fs::read_to_string(&self.data_path)?;
            let agents_data: AgentsData = serde_json::from_str(&content)?;
            *self.data.blocking_lock() = agents_data;
        }
        Ok(())
    }

    fn save_data_sync(&self) -> anyhow::Result<()> {
        let content = serde_json::to_string_pretty(&*self.data.blocking_lock())?;
        fs::write(&self.data_path, content)?;
        Ok(())
    }

    fn get_next_id_sync(&self) -> (i64, String) {
        let data = self.data.blocking_lock();
        let sequential_id = data
            .agents
            .iter()
            .map(|a| a.base_config.id)
            .max()
            .unwrap_or(0)
            + 1;
        let uuid = Uuid::new_v4().to_string();
        (sequential_id, uuid)
    }

    pub fn create_agent_sync(
        &self,
        name: &str,
        provider: &str,
        model: &str,
        preamble: &str,
        base_url: Option<&str>,
        tools: Option<&[String]>,
        category: AgentCategory,
        api_key: Option<&str>,
    ) -> anyhow::Result<i64> {
        let (id, uuid) = self.get_next_id_sync();
        let created_at = chrono::Utc::now().to_rfc3339();

        let base_config = AgentConfig {
            id,
            uuid,
            name: name.to_string(),
            provider: provider.to_string(),
            model: model.to_string(),
            preamble: preamble.to_string(),
            created_at,
            base_url: base_url.map(|s| s.to_string()),
            tools: tools.map(|t| t.to_vec()),
            api_key: api_key.map(|s| s.to_string()),
        };

        let agent = EnhancedAgentConfig {
            base_config,
            category,
            knowledge_base: None,
            vector_db: None,
        };

        let mut data = self.data.blocking_lock();
        data.agents.push(agent);
        drop(data);

        self.save_data_sync()?;
        Ok(id)
    }

    pub fn update_agent_knowledge_base_sync(
        &self,
        agent_id: i64,
        kb_config: KnowledgeBaseConfig,
    ) -> anyhow::Result<()> {
        let mut data = self.data.blocking_lock();
        let agent = data
            .agents
            .iter_mut()
            .find(|agent| agent.base_config.id == agent_id);

        if let Some(agent) = agent {
            agent.knowledge_base = Some(kb_config);
            drop(data);
            self.save_data_sync()?;
            Ok(())
        } else {
            Err(anyhow::anyhow!("Agent with id {} not found", agent_id))
        }
    }

    pub fn update_agent_vector_db_sync(
        &self,
        agent_id: i64,
        vector_db_config: VectorDbConfig,
    ) -> anyhow::Result<()> {
        let mut data = self.data.blocking_lock();
        let agent = data
            .agents
            .iter_mut()
            .find(|agent| agent.base_config.id == agent_id);

        if let Some(agent) = agent {
            agent.vector_db = Some(vector_db_config);
            drop(data);
            self.save_data_sync()?;
            Ok(())
        } else {
            Err(anyhow::anyhow!("Agent with id {} not found", agent_id))
        }
    }

    pub fn delete_agent_sync(&self, agent_id: i64) -> anyhow::Result<bool> {
        let mut data = self.data.blocking_lock();
        let initial_len = data.agents.len();
        data.agents.retain(|agent| agent.base_config.id != agent_id);
        let deleted = data.agents.len() < initial_len;
        drop(data);

        if deleted {
            self.save_data_sync()?;
        }
        Ok(deleted)
    }

    pub fn list_agents_sync(&self) -> anyhow::Result<Vec<EnhancedAgentConfig>> {
        let data = self.data.blocking_lock();
        let mut agents = data.agents.clone();
        agents.sort_by(|a, b| b.base_config.created_at.cmp(&a.base_config.created_at));
        Ok(agents)
    }

    pub fn get_agent_sync(&self, agent_id: i64) -> anyhow::Result<Option<EnhancedAgentConfig>> {
        let data = self.data.blocking_lock();
        Ok(data
            .agents
            .iter()
            .find(|agent| agent.base_config.id == agent_id)
            .cloned())
    }

    pub fn get_agent_by_category_sync(
        &self,
        category: &AgentCategory,
    ) -> anyhow::Result<Vec<EnhancedAgentConfig>> {
        let data = self.data.blocking_lock();
        let agents: Vec<EnhancedAgentConfig> = data
            .agents
            .iter()
            .filter(|agent| &agent.category == category)
            .cloned()
            .collect();
        Ok(agents)
    }

    // Async versions of the methods
    async fn load_data(&self) -> anyhow::Result<()> {
        if Path::new(&self.data_path).exists() {
            let content = fs::read_to_string(&self.data_path)?;
            let agents_data: AgentsData = serde_json::from_str(&content)?;
            *self.data.lock().await = agents_data;
        }
        Ok(())
    }

    async fn save_data(&self) -> anyhow::Result<()> {
        let content = serde_json::to_string_pretty(&*self.data.lock().await)?;
        fs::write(&self.data_path, content)?;
        Ok(())
    }

    async fn get_next_id(&self) -> (i64, String) {
        let data = self.data.lock().await;
        let sequential_id = data
            .agents
            .iter()
            .map(|a| a.base_config.id)
            .max()
            .unwrap_or(0)
            + 1;
        let uuid = Uuid::new_v4().to_string();
        (sequential_id, uuid)
    }

    pub async fn create_agent(
        &self,
        name: &str,
        provider: &str,
        model: &str,
        preamble: &str,
        base_url: Option<&str>,
        tools: Option<&[String]>,
        category: AgentCategory,
        api_key: Option<&str>,
    ) -> anyhow::Result<i64> {
        let (id, uuid) = self.get_next_id().await;
        let created_at = chrono::Utc::now().to_rfc3339();

        let base_config = AgentConfig {
            id,
            uuid,
            name: name.to_string(),
            provider: provider.to_string(),
            model: model.to_string(),
            preamble: preamble.to_string(),
            created_at,
            base_url: base_url.map(|s| s.to_string()),
            tools: tools.map(|t| t.to_vec()),
            api_key: api_key.map(|s| s.to_string()),
        };

        let agent = EnhancedAgentConfig {
            base_config,
            category,
            knowledge_base: None,
            vector_db: None,
        };

        let mut data = self.data.lock().await;
        data.agents.push(agent);
        drop(data);

        self.save_data().await?;
        Ok(id)
    }

    pub async fn update_agent_knowledge_base(
        &self,
        agent_id: i64,
        kb_config: KnowledgeBaseConfig,
    ) -> anyhow::Result<()> {
        let mut data = self.data.lock().await;
        let agent = data
            .agents
            .iter_mut()
            .find(|agent| agent.base_config.id == agent_id);

        if let Some(agent) = agent {
            agent.knowledge_base = Some(kb_config);
            drop(data);
            self.save_data().await?;
            Ok(())
        } else {
            Err(anyhow::anyhow!("Agent with id {} not found", agent_id))
        }
    }

    pub async fn update_agent_vector_db(
        &self,
        agent_id: i64,
        vector_db_config: VectorDbConfig,
    ) -> anyhow::Result<()> {
        let mut data = self.data.lock().await;
        let agent = data
            .agents
            .iter_mut()
            .find(|agent| agent.base_config.id == agent_id);

        if let Some(agent) = agent {
            agent.vector_db = Some(vector_db_config);
            drop(data);
            self.save_data().await?;
            Ok(())
        } else {
            Err(anyhow::anyhow!("Agent with id {} not found", agent_id))
        }
    }

    pub async fn delete_agent(&self, agent_id: i64) -> anyhow::Result<bool> {
        let mut data = self.data.lock().await;
        let initial_len = data.agents.len();
        data.agents.retain(|agent| agent.base_config.id != agent_id);
        let deleted = data.agents.len() < initial_len;
        drop(data);

        if deleted {
            self.save_data().await?;
        }
        Ok(deleted)
    }

    pub async fn list_agents(&self) -> anyhow::Result<Vec<EnhancedAgentConfig>> {
        let data = self.data.lock().await;
        let mut agents = data.agents.clone();
        agents.sort_by(|a, b| b.base_config.created_at.cmp(&a.base_config.created_at));
        Ok(agents)
    }

    pub async fn get_agent(&self, agent_id: i64) -> anyhow::Result<Option<EnhancedAgentConfig>> {
        let data = self.data.lock().await;
        Ok(data
            .agents
            .iter()
            .find(|agent| agent.base_config.id == agent_id)
            .cloned())
    }

    pub async fn get_agent_by_category(
        &self,
        category: &AgentCategory,
    ) -> anyhow::Result<Vec<EnhancedAgentConfig>> {
        let data = self.data.lock().await;
        let agents: Vec<EnhancedAgentConfig> = data
            .agents
            .iter()
            .filter(|agent| &agent.category == category)
            .cloned()
            .collect();
        Ok(agents)
    }
}
